"""
子通道分析程序主入口

用于运行整个子通道分析程序
"""

import argparse
import time
from subchannel_analysis.solvers.mass_momentum_energy import TwoFluidSolver
from subchannel_analysis.solvers.convergence import ConvergenceChecker
from subchannel_analysis.utils.plotting import ResultPlotter
from subchannel_analysis.utils.data_processing import DataProcessor
from subchannel_analysis.utils.data_logger import DataLogger
from subchannel_analysis import config as cfg

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='子通道分析程序')
    
    # 添加计算参数
    parser.add_argument('--n_timesteps', type=int, default=1000,
                      help='时间步数')
    parser.add_argument('--save_results', action='store_true',
                      help='是否保存结果')
    parser.add_argument('--plot_results', action='store_true',
                      help='是否绘制结果')
    parser.add_argument('--output_prefix', type=str, default='results',
                      help='输出文件前缀')
    
    return parser.parse_args()

def main():
    """主程序"""
    # 解析命令行参数
    args = parse_arguments()
    
    print("开始子通道分析计算...")
    print("-" * 50)
    
    # 记录开始时间
    start_time = time.time()
    
    # 创建数据记录器
    logger = DataLogger()
    logger.save_config()
    logger.log_message("开始计算...")
    
    # 创建求解器
    solver = TwoFluidSolver()
    
    # 创建收敛性检查器
    checker = ConvergenceChecker()
    
    # 运行模拟
    print("\n1. 运行数值模拟...")
    solver.run_simulation(args.n_timesteps)
    
    # 获取计算结果
    channels = solver.channels
    
    # 检查收敛性
    print("\n2. 检查计算收敛性...")
    converged, results = checker.check_convergence(channels, solver.channels)
    checker.print_convergence_status(results)
    
    # 处理结果
    print("\n3. 处理计算结果...")
    processor = DataProcessor()
    
    # 计算平均值
    averages = processor.calculate_average_values(channels)
    
    # 计算质量流量
    mass_flows = processor.calculate_mass_flow_rates(channels)
    
    # 计算能量平衡
    energy_balance = processor.calculate_energy_balance(channels)
    
    # 计算压降
    pressure_drops = processor.calculate_pressure_drop(channels)
    
    # 计算传热系数
    htc = processor.calculate_heat_transfer_coefficients(channels)
    
    # 计算统计数据
    stats = processor.calculate_statistics(channels)
    
    # 保存结果
    if args.save_results:
        print("\n4. 保存计算结果...")
        logger.save_field_data(channels)
        logger.save_results_summary({
            'averages': averages,
            'mass_flows': mass_flows,
            'energy_balance': energy_balance,
            'pressure_drops': pressure_drops,
            'heat_transfer_coefficients': htc,
            'statistics': stats
        })
        
    # 绘制结果
    if args.plot_results:
        print("\n5. 绘制计算结果...")
        plotter = ResultPlotter()
        
        # 绘制液相结果
        plotter.plot_liquid_phase_results(channels)
        
        # 绘制气相结果
        plotter.plot_vapor_phase_results(channels)
        
        # 绘制每个通道的轴向分布
        for channel_id in channels.keys():
            plotter.plot_axial_profiles(channels, channel_id)
            
        # 比较不同通道的温度分布
        plotter.plot_channel_comparison(channels, 'T_l')
        
        if args.save_results:
            plotter.save_results(channels, args.output_prefix)
    
    # 记录性能数据
    end_time = time.time()
    calculation_time = end_time - start_time
    logger.save_performance_data({
        'total_time': calculation_time,
        'simulation': calculation_time * 0.8,
        'post_processing': calculation_time * 0.2
    })
    
    # 打印计算时间
    print(f"\n计算总耗时: {calculation_time:.2f} 秒")
    
    # 打印主要结果
    print("\n计算结果摘要：")
    print("-" * 50)
    print("\n1. 平均值：")
    print(f"液相分数: {stats['alpha_l']['mean']:.4f} ± {stats['alpha_l']['std']:.4f}")
    print(f"液相温度: {stats['T_l']['mean']:.2f} ± {stats['T_l']['std']:.2f} K")
    print(f"壁面温度: {stats['T_w']['mean']:.2f} ± {stats['T_w']['std']:.2f} K")
    
    print("\n2. 压降：")
    for channel_id, drops in pressure_drops.items():
        print(f"通道 {channel_id}:")
        print(f"  总压降: {drops['total']/1e6:.2f} MPa")
        print(f"  重力压降: {drops['gravity']/1e6:.2f} MPa")
        print(f"  摩擦压降: {drops['friction']/1e6:.2f} MPa")
    
    print("\n3. 能量平衡：")
    total_imbalance = sum(bal['imbalance'] for bal in energy_balance.values())
    print(f"总能量不平衡: {total_imbalance/1e3:.2f} kW")
    
    print("\n计算完成！")
    print("-" * 50)
    
    # 记录计算完成
    logger.log_message("计算完成")

if __name__ == "__main__":
    main() 